#include "base/TokenBucket.h"
#include <gtest/gtest.h>
#include <thread>

// 测试TokenBucket类的基本功能
TEST(TokenBucketTest, BasicConsumption)
{
    // 创建一个TokenBucket实例，每秒产生10个令牌，最大突发大小为100
    TokenBucket<> bucket(10, 100);

    // 检查是否可以消耗1个令牌
    EXPECT_TRUE(bucket.consume(1)); // 应该返回true，因为桶是空的，并且有足够的时间来填充1个令牌

    // 检查是否可以消耗100个令牌
    EXPECT_TRUE(bucket.consume(100)); // 应该返回true，因为桶可以容纳100个令牌

    // 检查是否可以立即消耗101个令牌
    EXPECT_FALSE(bucket.consume(101)); // 应该返回false，因为桶中没有足够的令牌

    // 等待一段时间，然后检查是否可以消耗101个令牌
    // 这里需要等待的时间应该大于等于1个令牌的时间，即timePerToken_
    std::this_thread::sleep_for(bucket.timePerToken_ + std::chrono::milliseconds(1));
    EXPECT_TRUE(bucket.consume(101)); // 应该返回true，因为现在有足够的时间来填充额外的1个令牌
}

// 测试TokenBucket类的多线程消耗功能
TEST(TokenBucketTest, MultiThreadConsumption)
{
    // 创建一个TokenBucket实例，每秒产生足够多的令牌以供10个线程使用，最大突发大小为1000
    TokenBucket<> bucket(100, 1000);

    // 定义一个函数，用于线程中的令牌消耗
    auto consumeTokens = [&bucket](uint64_t tokens)
    {
        while (bucket.consume(tokens))
        {
            // 成功消耗tokens个令牌
        }
    };

    // 创建10个线程，每个线程尝试消耗10个令牌
    std::vector<std::thread> threads;
    for (int i = 0; i < 10; ++i)
    {
        threads.emplace_back(consumeTokens, 10);
    }

    // 等待所有线程完成
    for (auto &thread : threads)
    {
        thread.join();
    }

    // 检查是否所有线程都能成功消耗令牌
    // 这里可以根据实际情况添加更多的断言来验证线程消耗令牌的行为
}

int main(int argc, char **argv)
{
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}